package com.kancy.springcloudconfig.util.service;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
 * 密码持有类
 *
 * @author dongzhuming
 */
public class CipherHolder {

    private static final String CIPHER_INSTANCE_NAME = "AES/GCM/NoPadding";
    private static final String SECRET_KEY_ALGORITHM = "AES";
    private static final int GCM_TAG_LENGTH = 128;
    private static final String PROVIDER = "SunJCE";

    private static final int ENCRYPTION_MODE = 1;
    private static final int DECRYPTION_MODE = 2;

    private Map<String, byte[]> keys = new HashMap<>();
    private String defaultKeyId;

    /**
     * 添加密钥，并生成密码对，无defaultCipher时，自动填充
     *
     * @param id        密钥ID
     * @param key       密钥
     * @param isDefault 是否为默认密钥
     */
    synchronized void addKey(String id, byte[] key, boolean isDefault) {
        keys.put(id, key);
        if (Objects.isNull(defaultKeyId) || defaultKeyId.isEmpty() || isDefault) {
            defaultKeyId = id;
        }
    }

    synchronized void addKey(String id, byte[] key) {
        addKey(id, key, false);
    }

    Cipher getEncryptionCipher() {
        if (defaultKeyId == null) {
            throw new RuntimeException();
        }
        return prepareCipher(ENCRYPTION_MODE, defaultKeyId);
    }

    Cipher getEncryptionCipher(String keyId) {
        if (!keys.containsKey(keyId)) {
            throw new RuntimeException();
        }
        return prepareCipher(ENCRYPTION_MODE, keyId);
    }

    Cipher getDecryptionCipher() {
        if (defaultKeyId == null) {
            throw new RuntimeException();
        }
        return prepareCipher(DECRYPTION_MODE, defaultKeyId);
    }

    Cipher getDecryptionCipher(String keyId) {
        if (!keys.containsKey(keyId)) {
            throw new RuntimeException();
        }
        return prepareCipher(DECRYPTION_MODE, keyId);
    }

    public void populateSecretKey(String id, String key) {
        addKey(id, key.getBytes(), true);
    }

    private Cipher prepareCipher(int operationMode, String keyId) {
        byte[] key = keys.get(keyId);
        try {
            Cipher cipher = Cipher.getInstance(CIPHER_INSTANCE_NAME, PROVIDER);
            byte[] iv = generateIv(cipher, keyId);
            cipher.init(operationMode, new SecretKeySpec(key, SECRET_KEY_ALGORITHM), new GCMParameterSpec(GCM_TAG_LENGTH, iv));
            return cipher;
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | InvalidAlgorithmParameterException | NoSuchProviderException e) {
            throw new RuntimeException(e);
        }
    }

    private byte[] generateIv(Cipher cipher, String keyId) {
        byte[] iv = new byte[cipher.getBlockSize()];
        byte[] keyIdBytes = keyId.getBytes();
        for (int i = 0; i < Integer.min(keyIdBytes.length, iv.length); i++) {
            iv[i] = (byte) ((keyIdBytes[i] + i * 0xF) % 0xFF);
        }
        return iv;
    }
}
